#!/usr/bin/python
'''
Created on 2011-9-26

@author: coneagoe
'''
from optparse import OptionParser as OptionParser
import io
import base64
import time
from Crypto.Cipher import AES as AES
import os, random, struct


gOptions, gArgs = None, None

gEncodedKeyLen = 24

def do_option():
    global gOptions, gArgs

    parser = OptionParser("usage: %prog [options] file")
    parser.add_option(r'-k', r'--key', dest = 'key', help = r'key')
    parser.add_option(r'-e', r'--encode', action = "store_true", dest = 'encode', help = r'encode file')
    parser.add_option(r'-d', r'--decode', action = "store_true", dest = 'decode', help = r'decode file')
    parser.add_option(r'-r', r'--rename', action = "store_true", dest = 'rename', help = r'rename file')
    gOptions, gArgs = parser.parse_args()


def encrypt_file(in_filename, out_filename = None, chunksize = 64 * 1024):
    """ Encrypts a file using AES (CBC mode) with the
        given key.

        key:
            The encryption key - a string that must be
            either 16, 24 or 32 bytes long. Longer keys
            are more secure.

        in_filename:
            Name of the input file

        out_filename:
            If None, '<in_filename>.enc' will be used.

        chunksize:
            Sets the size of the chunk which the function
            uses to read and encrypt the file. Larger chunk
            sizes can be faster for some files and machines.
            chunksize must be divisible by 16.
    """
    global gOptions

    if not out_filename:
        out_filename = in_filename + '.tmp'

    if not gOptions.key:
        key = base64.b64encode(str(int(time.mktime(time.localtime()))))
    else:
        key = gOptions.key

    iv = ''.join(chr(random.randint(0, 0xFF)) for i in range(16))
    encryptor = AES.new(key, AES.MODE_CBC, iv)
    filesize = os.path.getsize(in_filename)

    with open(in_filename, 'rb') as infile:
        with open(out_filename, 'wb') as outfile:
            outfile.write(struct.pack('<Q', filesize))

            outfile.write(struct.pack('<16s', key))

            outfile.write(iv)

            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                elif len(chunk) % 16 != 0:
                    chunk += ' ' * (16 - len(chunk) % 16)

                outfile.write(encryptor.encrypt(chunk))

    if gOptions.rename:
        os.rename(out_filename, in_filename)


def decrypt_file(in_filename, out_filename = None, chunksize = 64 * 1024):
    """ Decrypts a file using AES (CBC mode) with the
        given key. Parameters are similar to encrypt_file,
        with one difference: out_filename, if not supplied
        will be in_filename without its last extension
        (i.e. if in_filename is 'aaa.zip.enc' then
        out_filename will be 'aaa.zip')
    """
    if not out_filename:
        out_filename = in_filename + '.tmp'

    with open(in_filename, 'rb') as infile:
        origsize = struct.unpack('<Q', infile.read(struct.calcsize('Q')))[0]

        if not gOptions.key:
            key = struct.unpack('<16s', infile.read(16))[0]
        else:
            key = gOptions.key

        iv = infile.read(16)
        decryptor = AES.new(key, AES.MODE_CBC, iv)

        with open(out_filename, 'wb') as outfile:
            while True:
                chunk = infile.read(chunksize)
                if len(chunk) == 0:
                    break
                outfile.write(decryptor.decrypt(chunk))

            outfile.truncate(origsize)

    if gOptions.rename:
        os.rename(out_filename, in_filename)


if __name__ == '__main__':
    do_option()

    if gOptions.encode:
        for input_file in gArgs:
            encrypt_file(input_file)
    elif gOptions.decode:
        for input_file in gArgs:
            decrypt_file(input_file)
